// SPDX-License-Identifier: Apache-2.0
import { illegalArgs } from "@thi.ng/errors/illegal-arguments";
import { BitInputStream } from "./input.js";

const DEFAULT_BUF_SIZE = 0x10;
const U32 = 0x1_0000_0000;

export class BitOutputStream {
	buffer: Uint8Array;
	protected start: number;
	protected pos!: number;
	protected bit!: number;
	protected bitPos!: number;

	constructor(buffer?: number | Uint8Array, offset = 0) {
		this.buffer =
			typeof buffer === "undefined"
				? new Uint8Array(DEFAULT_BUF_SIZE)
				: typeof buffer === "number"
				? new Uint8Array(buffer)
				: buffer;
		this.start = offset;
		this.seek(offset);
		this.buffer[this.pos] &= ~((1 << this.bit) - 1);
	}

	get position() {
		return this.bitPos;
	}

	seek(pos: number): BitOutputStream {
		if (pos < this.start || pos >= this.buffer.length << 3) {
			illegalArgs(`seek pos out of bounds: ${pos}`);
		}
		this.pos = pos >>> 3;
		this.bit = 8 - (pos & 0x7);
		this.bitPos = pos;
		return this;
	}

	bytes() {
		return this.buffer.slice(0, this.pos + (this.bit & 7 ? 1 : 0));
	}

	reader(from = 0) {
		return new BitInputStream(this.buffer, from, this.position);
	}

	write(x: number, wordSize = 1) {
		if (wordSize > 32) {
			let hi = Math.floor(x / U32);
			this.write(hi, wordSize - 32);
			this.write(x - hi * U32, 32);
		} else if (wordSize > 8) {
			let n = wordSize & -8;
			let msb = wordSize - n;
			if (msb > 0) {
				this._write(x >>> n, msb);
			}
			n -= 8;
			while (n >= 0) {
				this._write(x >>> n, 8);
				n -= 8;
			}
		} else {
			this._write(x, wordSize);
		}
		return this;
	}

	writeWords(input: Iterable<number>, wordSize = 8) {
		let iter = input[Symbol.iterator]();
		let v: IteratorResult<number>;
		while (((v = iter.next()), !v.done)) {
			this.write(v.value, wordSize);
		}
	}

	writeBit(x: number) {
		this.bit--;
		this.buffer[this.pos] =
			(this.buffer[this.pos] & ~(1 << this.bit)) | (x << this.bit);
		if (this.bit === 0) {
			this.ensureSize();
			//this.buffer[this.pos] = 0;
			this.bit = 8;
		}
		this.bitPos++;
		return this;
	}

	protected _write(x: number, wordSize: number) {
		x &= (1 << wordSize) - 1;
		let buf = this.buffer;
		let pos = this.pos;
		let bit = this.bit;
		let b = bit - wordSize;
		let m = bit < 8 ? ~((1 << bit) - 1) : 0;
		if (b >= 0) {
			m |= (1 << b) - 1;
			buf[pos] = (buf[pos] & m) | ((x << b) & ~m);
			if (b === 0) {
				this.ensureSize();
				this.bit = 8;
			} else {
				this.bit = b;
			}
		} else {
			this.bit = bit = 8 + b;
			buf[pos] = (buf[pos] & m) | ((x >>> -b) & ~m);
			this.ensureSize();
			this.buffer[this.pos] =
				(this.buffer[this.pos] & ((1 << bit) - 1)) |
				((x << bit) & 0xff);
		}
		this.bitPos += wordSize;
		return this;
	}

	protected ensureSize() {
		if (++this.pos === this.buffer.length) {
			let b = new Uint8Array(this.buffer.length << 1);
			b.set(this.buffer);
			this.buffer = b;
		}
	}
}
